# -*- coding:utf-8 -*-

from sklearn.datasets import load_iris,load_boston
from sklearn.feature_selection import VarianceThreshold
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.linear_model import LinearRegression,SGDRegressor,Ridge
from sklearn.metrics import mean_squared_error
import joblib
import pandas as pd
import numpy as np
import sys
import pypinyin
import math

sys.path.append("../")
from frameworks.utils.PadasExcelUtil import *
import re
import nums_from_string
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV
from sklearn.decomposition import PCA
from scipy.stats import pearsonr
import matplotlib.pyplot as plt
from services.DayKlineService import *
import warnings
warnings.filterwarnings('ignore')

def main(code):
    # 1）获取数据
    codeService = DayKlineService();
    data = codeService.getData(code);

    origDf = pd.DataFrame(data)
    print("特征数量：\n", origDf.shape)
    origDf["Close"] = origDf["close"]
    origDf["High"] = origDf["high"]
    origDf["Low"] = origDf["low"]
    origDf["Open"] = origDf["open"]
    origDf["Volume"] = origDf["volume"]
    origDf["Date"] = origDf["trade_date"]
    print(origDf)

    df = origDf[['Close', 'High', 'Low', 'Open', 'Volume']]
    featureData = df[['Open', 'High', 'Volume', 'Low']]
    # 划分特征值和目标值
    feature = featureData.values
    target = np.array(df['Close'])

    # 划分训练集，测试集
    feature_train, feature_test, target_train, target_test = train_test_split(feature, target, test_size=0.05)
    pridectedDays = int(math.ceil(0.05 * len(origDf)))  # 预测天数
    lrTool = LinearRegression()
    lrTool.fit(feature_train, target_train)  # 训练
    # 用测试集预测结果
    predictByTest = lrTool.predict(feature_test)

    # 组装数据
    index = 0
    # 在前95%的交易日中，设置预测结果和收盘价一致
    while index < len(origDf) - pridectedDays:
        df.loc[index, 'predictedVal'] = origDf.loc[index, 'Close']
        df.loc[index, 'Date'] = origDf.loc[index, 'Date']
        index = index + 1
    predictedCnt = 0
    # 在后5%的交易日中，用测试集推算预测股价
    while predictedCnt < pridectedDays:
        df.loc[index, 'predictedVal'] = predictByTest[predictedCnt]
        df.loc[index, 'Date'] = origDf.loc[index, 'Date']
        predictedCnt = predictedCnt + 1
        index = index + 1

    plt.figure()
    df['predictedVal'].plot(color="red", label='predicted Data')
    df['Close'].plot(color="blue", label='Real Data')
    plt.legend(loc='best')  # 绘制图例
    # 设置x坐标的标签
    major_index = df.index[df.index % 10 == 0]
    major_xtics = df['Date'][df.index % 10 == 0]
    plt.xticks(major_index, major_xtics)
    plt.setp(plt.gca().get_xticklabels(), rotation=30)
    # 带网格线，且设置了网格样式
    plt.grid(linestyle='-.')
    plt.show()

if __name__ == "__main__":
    main("SH600010")